package com.zh.stargraph2.core.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.zh.stargraph2.core.dto.request.Text2ImageReqDto;
import com.zh.stargraph2.core.dto.respone.Text2ImageResDto;
import com.zh.stargraph2.core.pojo.ComfyuiModel;
import com.zh.stargraph2.core.pojo.ComfyuiRequestDto;
import com.zh.stargraph2.core.pojo.ComfyuiTask;
import com.zh.stargraph2.core.service.FreemarkerService;
import com.zh.stargraph2.core.service.OllamaService;
import com.zh.stargraph2.core.service.RedisService;
import com.zh.stargraph2.core.service.Text2ImageService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Service
public class Text2ImageServiceImpl implements Text2ImageService {

    public static final String COMFYUI_CLIENT_ID = "star-graph";

    @Autowired
    private OllamaService ollamaService;

    @Autowired
    private FreemarkerService freemarkerService;

    @Autowired
    RedisService redisService;

    @Override
    public Text2ImageResDto prompt(Text2ImageReqDto text2ImageReqDto){
        try {
            ComfyuiTask comfyuiTask = getComfyuiTask(text2ImageReqDto);
            comfyuiTask = redisService.addQueueTask(comfyuiTask);
            Text2ImageResDto text2ImageResDto = new Text2ImageResDto();
            text2ImageResDto.setPid(comfyuiTask.getId());
            text2ImageResDto.setQueueIndex(comfyuiTask.getIndex());
            return text2ImageResDto;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public ComfyuiTask getComfyuiTask(Text2ImageReqDto text2ImageReqDto) throws Exception {
        // 1. 将text2ImageReqDto转换成ComfyuiModel
        // 基本参数封装
        ComfyuiModel model = BeanUtil.toBean(text2ImageReqDto, ComfyuiModel.class);
        model.setWidth(text2ImageReqDto.width());
        model.setHeight(text2ImageReqDto.height());
        model.setModelName(text2ImageReqDto.modelName());
        model.setSamplerName(text2ImageReqDto.samplerName());
        model.setScheduler(text2ImageReqDto.scheduler());
        // 封装提示词与负向提示词
        model.setPropmt(String.format("(8k, best quality, masterpiece),(high detailed skin),%s", ollamaService.translate(model.getPropmt())));
        if (StrUtil.isNotEmpty(model.getReverse())) {
            model.setReverse(ollamaService.translate(model.getReverse()));
        }
        model.setReverse(String.format("%s,bad face,naked,bad finger,bad arm,bad leg,bad eye", model.getReverse()));

        // 2.利用ComfyuiModel结合模板使用Freemarker渲染出comfyui json
        String prompt = freemarkerService.renderText2Image(model);

        // 3.构建出Comfyui生图需要的对象
        ComfyuiRequestDto dto = new ComfyuiRequestDto(COMFYUI_CLIENT_ID, JSON.parseObject(prompt));

        // 4.构建出ComfyuiTask（Redis任务对象）
        return new ComfyuiTask(text2ImageReqDto.getClientId(), dto);
    }


}
